#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
文件名称: maples_merge_preprocess.py
功能描述:
    合并 MAPLES-DR 与 Messidor 标注信息，
    同时进行 mask 裁剪与日志记录。
"""

import os
import json
import csv
import argparse
from tqdm import tqdm
from data_preprocess.utils.crop import crop_resize_save_with_info
from data_preprocess.utils.stat import save_classification_stats

def parse_drme(dr, me):
    """将 DR/ME 标注转换为文本描述"""
    dr_map = {
        "R0": "normal",
        "R1": "Mild Diabetic Retinopathy",
        "R2": "Moderate Diabetic Retinopathy",
        "R3": "Severe Diabetic Retinopathy",
        "R4A": "Proliferative Diabetic Retinopathy",
        "R4S": "Stable treated proliferative Diabetic Retinopathy",
        "R6": "Ungradable quality"
    }
    me_map = {
        "M0": "normal",
        "M1": "Referable Diabetic Macular Edema",
        "M2": "Referable Diabetic Macular Edema",
        "M6": "Ungradable quality"
    }

    # 跳过质量不足的标注
    if me == "M6" or dr == "R6":
        return None

    dr_text = dr_map.get(dr.strip(), "normal")
    me_text = me_map.get(me.strip(), "normal")

    if dr_text == "normal" and me_text == "normal":
        return "normal"
    elif dr_text != "normal" and me_text == "normal":
        return dr_text
    elif dr_text == "normal" and me_text != "normal":
        return me_text
    else:
        return f"{dr_text}, {me_text}"


def load_old_annotations(processed_root):
    """加载 Messidor 已处理结果"""
    anno_path = os.path.join(processed_root, "annotations.json")
    if not os.path.exists(anno_path):
        raise FileNotFoundError(f"未找到 Messidor 标注文件: {anno_path}")
    with open(anno_path, "r", encoding="utf-8") as f:
        annotations = json.load(f)

    id2key = {}
    for key, val in annotations.items():
        image_id = os.path.splitext(os.path.basename(val["original_path"]))[0]
        id2key[image_id] = key
    return annotations, id2key


def load_new_diagnosis(MAPLES_DR_root):
    """加载 MAPLES-DR 的 diagnosis.csv"""
    id2diag = {}
    for subset in ["train", "test"]:
        csv_path = os.path.join(MAPLES_DR_root, "MAPLES-DR", subset, "diagnosis.csv")
        if not os.path.exists(csv_path):
            print(f"[WARN] 未找到 {subset} 的 CSV 文件: {csv_path}")
            continue
        with open(csv_path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                name, dr, me = row["name"].strip(), row["DR"].strip(), row["ME"].strip()
                parsed = parse_drme(dr, me)
                if parsed is not None:
                    id2diag[name] = parsed
    return id2diag


def merge_classification(old_annos, id2key, id2diag, mismatch_log, missing_log):
    """合并分类标签"""
    for image_id, new_text in tqdm(id2diag.items(), desc="🔁 合并分类标签"):
        if image_id not in id2key:
            missing_log.append(image_id)
            continue

        key = id2key[image_id]
        old_diag = old_annos[key].get("diagnosis", {}).get("classification", {})
        old_text = old_diag.get("text", "normal")

        merged_text = new_text
        if old_text != new_text and new_text not in old_text and old_text not in new_text:
            # 记录冲突
            mismatch_log.append({
                "id": image_id,
                "old": old_text,
                "new": new_text
            })

        if "diagnosis" not in old_annos[key]:
            old_annos[key]["diagnosis"] = {}
        if "classification" not in old_annos[key]["diagnosis"]:
            old_annos[key]["diagnosis"]["classification"] = {}

        # ✅ 冲突时优先使用新标签
        old_annos[key]["diagnosis"]["classification"]["text"] = merged_text
    return old_annos

def merge_segmentation(old_annos, id2key, MAPLES_DR_root, save_root, resize=(512, 512)):
    """合并分割标注并处理命名映射"""
    # 文件夹类别名（数据集文件夹）
    seg_categories = [
        "Drusens", "Exudates", "Hemorrhages", "Macula",
        "Microaneurysms", "Neovascularization", "OpticCup",
        "OpticDisc", "CottonWoolSpots", "Vessels"
    ]

    # ✅ 文件夹名 -> 标准化 key 的映射（统一小写，有空格）
    folder_to_key = {
        "Drusens": "drusens",
        "Exudates": "exudates",
        "Hemorrhages": "hemorrhages",
        "Macula": "macula",
        "Microaneurysms": "microaneurysms",
        "Neovascularization": "neovascularization",
        "OpticCup": "optic cup",
        "OpticDisc": "optic disc",
        "CottonWoolSpots": "cotton-wool spots",  # 注意这里用中划线
        "Vessels": "vessels"
    }

    for subset in ["train", "test"]:
        subset_dir = os.path.join(MAPLES_DR_root, "MAPLES-DR", subset)
        for cat in seg_categories:
            seg_dir = os.path.join(subset_dir, cat)
            if not os.path.isdir(seg_dir):
                continue

            # 使用 key 映射确定保存目录（例如 optic cup → optic_cup）
            key_name = folder_to_key[cat]
            save_dir = os.path.join(save_root, key_name.replace(" ", "_"))
            os.makedirs(save_dir, exist_ok=True)

            files = [
                f for f in os.listdir(seg_dir)
                if f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff"))
            ]

            for fname in tqdm(files, desc=f"🎨 处理 {subset} 分割 - {cat}", leave=False, unit="masks"):
                image_id = os.path.splitext(fname)[0]
                if image_id not in id2key:
                    # 未在旧数据中匹配到原图，跳过
                    continue

                key = id2key[image_id]
                crop_info = old_annos.get(key, {}).get("crop_info")
                if not crop_info:
                    raise ValueError(f"❌ no crop info for {image_id}")

                image_name = old_annos[key]["image_name"]
                mask_ori_path = os.path.join(seg_dir, fname)
                mask_save_path = os.path.join(save_dir, image_name)

                try:
                    crop_resize_save_with_info(
                        image_path=mask_ori_path,
                        save_path=mask_save_path,
                        info=crop_info,
                        resize=resize
                    )
                except Exception as e:
                    raise RuntimeError(
                        f"处理 mask 失败: subset={subset}, cat={cat}, fname={fname}, "
                        f"key={key}, image_name={image_name}, 错误: {e}"
                    )

                # ✅ 存储时用统一 key 名（而不是 cat.lower()）
                diag = old_annos[key].setdefault("diagnosis", {})
                seg_bucket = diag.setdefault("segmentation", {})
                seg_bucket[key_name] = mask_save_path

    return old_annos


def main(args):
    os.makedirs("experiments", exist_ok=True)

    mismatch_log, missing_log = [], []

    print("📂 加载 Messidor annotations ...")
    old_annos, id2key = load_old_annotations(args.Messidor_processed_root)

    print("📄 读取 MAPLES-DR diagnosis 信息 ...")
    id2diag = load_new_diagnosis(args.MAPLES_DR_root)

    print("🔁 合并分类标签 ...")
    old_annos = merge_classification(old_annos, id2key, id2diag, mismatch_log, missing_log)

    print("🎨 合并分割标注并执行裁剪 ...")
    old_annos = merge_segmentation(old_annos, id2key, args.MAPLES_DR_root, args.save_root, resize=(512, 512))

    os.makedirs(args.save_root, exist_ok=True)
    save_path = os.path.join(args.save_root, "annotations.json")
    with open(save_path, "w", encoding="utf-8") as f:
        json.dump(old_annos, f, indent=4, ensure_ascii=False)
    print(f"✅ 合并完成，结果已保存至: {save_path}")

    mismatch_label_path = './experiments/maplesdr_log/mismatch_cls.json'
    missing_path = './experiments/maplesdr_log/missing.json'
    # 写入日志文件
    os.makedirs('./experiments/maplesdr_log/',exist_ok=True)
    with open(mismatch_label_path, "w", encoding="utf-8") as f:
        json.dump(mismatch_log, f, indent=4, ensure_ascii=False)
    with open(missing_path, "w", encoding="utf-8") as f:
        json.dump(missing_log, f, indent=4, ensure_ascii=False)

    # stat
    save_classification_stats(old_annos,'./experiments/stat/maplesdr.json')
    print("📄 冲突记录已保存至 "+mismatch_label_path)
    print("📄 无匹配原图记录已保存至 "+missing_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="合并 MAPLES-DR 与 Messidor 标注信息（含日志输出）")
    parser.add_argument("--Messidor_processed_root", type=str,
                        default="/home/zhangpinglu/data0/gy/Dataset/public_processed/Messidor",
                        help="Messidor 预处理后数据路径（含 annotations.json）")
    parser.add_argument("--MAPLES_DR_root", type=str,
                        default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/MAPLES-DR",
                        help="MAPLES-DR 原始数据路径")
    parser.add_argument("--save_root", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/MAPLES-DR",
                        help="合并后结果保存路径（含裁剪后mask）")
    args = parser.parse_args()
    main(args)
